import torch
import game_layer_cpp
from net import ReversiNet
from tqdm import tqdm
from torch.distributions import Categorical

ITERS = 100
BATCH_SIZE = 256
BOARD_SIZE = 15

@torch.no_grad()
def eval(model, model_baseline_path='data/gomoku_231024.pth', device='cuda'):
    model_baseline = torch.jit.load(model_baseline_path).to(device)
    board = torch.zeros((BATCH_SIZE, 2, BOARD_SIZE, BOARD_SIZE), device=device, dtype=torch.bool)
    board[:BATCH_SIZE//2, 1, 7, 7] = 1

    pbar = tqdm(range(ITERS + 1))
    win_model = 0
    win_baseline = 0
    win_total_stone = 0
    for iter in pbar:
        p: torch.Tensor = model(board.float())[:, 0].flatten(1)
        invalid_mask = board.any(1).flatten(1)
        p = p.masked_fill(invalid_mask, float('-inf'))
        action = Categorical(logits=p / 0.001).sample()

        # step game
        win = game_layer_cpp.check_win_cuda(board, action)
        win_model += win.sum().item()
        win_total_stone += board[win].sum().item()

        board[win | board.any(1).flatten(1).all(1)] = False
        board = torch.flip(board, (1,))

        # run baseline
        p: torch.Tensor = model_baseline(board.float()).flatten(1)
        invalid_mask = board.any(1).flatten(1)
        p = p.masked_fill(invalid_mask, float('-inf'))
        action = Categorical(logits=p / 0.05).sample()

        # step game
        win = game_layer_cpp.check_win_cuda(board, action)
        win_baseline += win.sum().item()
        win_total_stone += board[win].sum().item()

        board[win | board.any(1).flatten(1).all(1)] = False
        board = torch.flip(board, (1,))
        win_cnt = win_baseline + win_model
        if win_cnt:
            win_rate = win_model / win_cnt
            pbar.set_description_str(f'{win_model}/{win_baseline}, {win_rate:.3f}')
    if win_cnt:
        return win_model / win_cnt, win_total_stone / win_cnt
    return None, None

@torch.no_grad()
def demo(model, model_baseline_path='data/gomoku_230201.pth', device='cuda'):
    # model_baseline = torch.jit.load(model_baseline_path).to(device)
    model_baseline = model
    board = torch.zeros((1, 2, BOARD_SIZE, BOARD_SIZE), device=device, dtype=torch.bool)

    pbar = tqdm(range(ITERS + 1))
    win_model = 0
    win_baseline = 0
    win_total_stone = 0
    for iter in pbar:
        p: torch.Tensor = model(board.float()).flatten(1)
        invalid_mask = board.any(1).flatten(1)
        p = p.masked_fill(invalid_mask, float('-inf'))
        action = Categorical(logits=p / 0.001).sample()

        # step game
        win = game_layer_cpp.check_win_cuda(board, action)
        win_model += win.sum().item()
        win_total_stone += board[win].sum().item()

        board[win | board.any(1).flatten(1).all(1)] = False
        board = torch.flip(board, (1,))
        # print(board[0, 0].int() + board[0, 1].int()*2)
        print(1-p.max().sigmoid())

        # run baseline
        p: torch.Tensor = model_baseline(board.float()).flatten(1)
        invalid_mask = board.any(1).flatten(1)
        p = p.masked_fill(invalid_mask, float('-inf'))
        action = Categorical(logits=p / 0.05).sample()

        # step game
        win = game_layer_cpp.check_win_cuda(board, action)
        win_baseline += win.sum().item()
        win_total_stone += board[win].sum().item()

        board[win | board.any(1).flatten(1).all(1)] = False
        board = torch.flip(board, (1,))
        # print(board[0, 0].int()*2 + board[0, 1].int())
        print(p.max().sigmoid())

        win_cnt = win_baseline + win_model
        if win_cnt:
            win_rate = win_model / win_cnt
            pbar.set_description_str(f'{win_model}/{win_baseline}, {win_rate:.3f}')
    if win_cnt:
        return win_model / win_cnt, win_total_stone / win_cnt
    return None, None

if __name__ == '__main__':
    from exps.run2.net import ReversiNet
    device = 'cuda'
    model = ReversiNet(train=True).to(device).eval()
    model.load_state_dict(torch.load('exps/run2/ema0020.pth', map_location='cpu'))
    # model = torch.nn.DataParallel(model)
    eval(model, device=device)
